import asyncio
import aiohttp
import time
import json
from concurrent.futures import ThreadPoolExecutor
import statistics

class LoadTester:
    def __init__(self, base_url="http://localhost:5000", max_concurrent=10000):
        self.base_url = base_url
        self.max_concurrent = max_concurrent
        self.results = []
        
    async def single_request(self, session, endpoint, method="GET", data=None):
        """单个请求测试"""
        start_time = time.time()
        try:
            if method == "GET":
                async with session.get(f"{self.base_url}{endpoint}") as response:
                    status = response.status
                    content = await response.text()
            elif method == "POST":
                async with session.post(f"{self.base_url}{endpoint}", json=data) as response:
                    status = response.status
                    content = await response.text()
            
            end_time = time.time()
            return {
                'status': status,
                'response_time': end_time - start_time,
                'success': status == 200,
                'endpoint': endpoint
            }
        except Exception as e:
            end_time = time.time()
            return {
                'status': 0,
                'response_time': end_time - start_time,
                'success': False,
                'error': str(e),
                'endpoint': endpoint
            }
    
    async def concurrent_test(self, endpoint, concurrent_users, duration=60):
        """并发测试"""
        print(f"开始测试 {endpoint}，并发用户数: {concurrent_users}，持续时间: {duration}秒")
        
        connector = aiohttp.TCPConnector(limit=concurrent_users * 2)
        timeout = aiohttp.ClientTimeout(total=30)
        
        async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
            start_time = time.time()
            tasks = []
            
            while time.time() - start_time < duration:
                # 创建并发任务
                for _ in range(min(concurrent_users, 1000)):  # 批次控制
                    task = asyncio.create_task(
                        self.single_request(session, endpoint)
                    )
                    tasks.append(task)
                
                # 等待一批完成
                if len(tasks) >= 1000:
                    results = await asyncio.gather(*tasks, return_exceptions=True)
                    self.results.extend([r for r in results if isinstance(r, dict)])
                    tasks = []
                    
                await asyncio.sleep(0.1)  # 短暂休息
            
            # 处理剩余任务
            if tasks:
                results = await asyncio.gather(*tasks, return_exceptions=True)
                self.results.extend([r for r in results if isinstance(r, dict)])
    
    def analyze_results(self):
        """分析测试结果"""
        if not self.results:
            print("没有测试结果")
            return
        
        successful = [r for r in self.results if r['success']]
        failed = [r for r in self.results if not r['success']]
        
        response_times = [r['response_time'] for r in successful]
        
        print("\n=== 测试结果分析 ===")
        print(f"总请求数: {len(self.results)}")
        print(f"成功请求: {len(successful)} ({len(successful)/len(self.results)*100:.2f}%)")
        print(f"失败请求: {len(failed)} ({len(failed)/len(self.results)*100:.2f}%)")
        
        if response_times:
            print(f"平均响应时间: {statistics.mean(response_times):.3f}秒")
            print(f"中位数响应时间: {statistics.median(response_times):.3f}秒")
            print(f"95%响应时间: {sorted(response_times)[int(len(response_times)*0.95)]:.3f}秒")
            print(f"最大响应时间: {max(response_times):.3f}秒")
            print(f"最小响应时间: {min(response_times):.3f}秒")
        
        # 错误统计
        if failed:
            error_types = {}
            for f in failed:
                error = f.get('error', f'HTTP {f["status"]}')
                error_types[error] = error_types.get(error, 0) + 1
            
            print("\n错误类型统计:")
            for error, count in error_types.items():
                print(f"  {error}: {count}")

async def run_load_test():
    """运行负载测试"""
    tester = LoadTester()
    
    # 测试场景
    test_scenarios = [
        # 基础功能测试
        {'endpoint': '/', 'users': 1000, 'duration': 30},
        {'endpoint': '/project_groups', 'users': 2000, 'duration': 30},
        {'endpoint': '/data', 'users': 5000, 'duration': 60},
        
        # 高并发测试
        {'endpoint': '/', 'users': 10000, 'duration': 60},
        {'endpoint': '/summary', 'users': 20000, 'duration': 120},
        
        # 极限测试
        {'endpoint': '/', 'users': 50000, 'duration': 180},
        {'endpoint': '/project_groups', 'users': 100000, 'duration': 300},
    ]
    
    for scenario in test_scenarios:
        tester.results = []  # 重置结果
        await tester.concurrent_test(
            scenario['endpoint'], 
            scenario['users'], 
            scenario['duration']
        )
        print(f"\n--- {scenario['endpoint']} 测试完成 ---")
        tester.analyze_results()
        print("\n" + "="*50 + "\n")
        
        # 测试间隔
        await asyncio.sleep(10)

if __name__ == "__main__":
    asyncio.run(run_load_test())